# 此脚本用于生成变化的反硝化速率常数和观测的硝态氮浓度值
import pandas as pd
import numpy as np
from plotstyles import figure
from plotstyles import fonts
from matplotlib import pyplot as plt

plt.rcParams["font.size"] = 8


def function_temp(temperature, theta, critical_temp):
    if temperature >= critical_temp:
        return np.power(theta, temperature - 20)
    else:
        return 0


def function_oxygen(oxygen, optimal_oxygen, critical_oxygen, beta):
    if oxygen < optimal_oxygen:
        return 1
    elif oxygen > critical_oxygen:
        return 0
    else:
        return (critical_oxygen - oxygen) / (critical_oxygen - optimal_oxygen) + (
            np.exp(beta) - np.e
        ) * (oxygen - optimal_oxygen)


def add_noise(true_value, relative_error):
    return true_value * (
        1 + np.random.normal(0, (relative_error / 3) ** 2, true_value.shape)
    )


def read_data(filename):
    data = pd.read_excel(filename)
    oxygen = data["DO"]
    temperature = data["T"]
    return oxygen, temperature


def calc_parameter_k(k20, T, DO, theta_deni, T_cdeni, c_oxo_deni, c_oxc_deni, beta):
    k = []
    for t, do in zip(T, DO):
        k.append(
            k20
            * function_temp(t, theta_deni, T_cdeni)
            * function_oxygen(do, c_oxo_deni, c_oxc_deni, beta)
        )
    return np.array(k)


def get_true_nitrate(k, c0, dt=1):
    """
    k为所有时刻的反应速率常数，c0为起始浓度, dt为时间间隔
    """
    c = [c0]
    for kt in k:
        c.append(c[-1] * np.exp(-kt * dt))
    return np.array(c[:-1])


if __name__ == "__main__":
    np.random.seed(0)
    DO, T = read_data("DataInit/TemperatureAndDO.xlsx")
    t = np.arange(0, DO.shape[0], 1)

    # 数据生成
    k20 = 0.0038
    theta_deni = 1.11
    T_cdeni = -2
    c_oxo_deni = 3.0
    c_oxc_deni = 10.0
    beta = 1
    c0 = 1
    re = 0.05
    k = calc_parameter_k(k20, T, DO, theta_deni, T_cdeni, c_oxo_deni, c_oxc_deni, beta)
    c_true = get_true_nitrate(k, c0, 1)
    c_obs = add_noise(c_true, re)

    dataout = pd.DataFrame()
    dataout['t'] = t
    dataout['k'] = k
    dataout['c_true'] = c_true
    dataout['c_obs'] = c_obs
    dataout.to_excel('DataGen/MiyunParameterTracker.xlsx', index=False)

    # 图形绘制
    fig = figure.Figure(14, 8)
    ax_do = fig.add_axes_cm("dot", 1, 0.5, 5, 3, "left upper")
    ax_do.set_title("   $(\mathbf{a})$", loc="left", y=0.8, pad=0)
    ax_do.set_ylim(6, 12)
    ax_do.set_ylabel("DO (mg/L)", color="b")
    ax_t = ax_do.twinx()
    ax_t.set_xlim(-10, 800)
    ax_t.set_ylim(0, 30)
    ax_t.set_ylabel("T ($^\circ C$)", color="r")
    ax_t.xaxis.set_visible(False)
    ax_do.xaxis.set_visible(False)

    ax_k = fig.add_axes_cm("k", 1, 4, 5, 3, "left upper")
    ax_k.set_title("   $(\mathbf{b})$", loc="left", y=0.8, pad=0)
    ax_k.ticklabel_format(style="sci", scilimits=(-3, 1), axis="y", useMathText=True)
    ax_k.set_ylabel("k (1/d)")
    ax_k.set_xlim(-10, 800)
    ax_k.set_ylim(0, 0.003)
    ax_k.set_xlabel("Time (d)")

    ax_cnitrate = fig.add_axes_cm("nitrate", 7.6, 0.5, 6, 6.5, "left upper")
    ax_cnitrate.set_title("    $(\mathbf{c})$", loc="left", y=0.9, pad=0)
    ax_cnitrate.set_xlim(-10, 800)
    ax_cnitrate.set_ylim(0.5, 1.0)
    ax_cnitrate.set_ylabel("Nitrate (mg/L)")
    ax_cnitrate.set_xlabel("Time (d)")
    miniax = ax_cnitrate.inset_axes([0.6, 0.6, 0.4, 0.4])
    miniax.set_xlim(630, 645)
    miniax.set_ylim(0.555, 0.565)
    miniax.xaxis.set_visible(False)
    miniax.yaxis.set_visible(False)

    ax_do.plot(t, DO, color="b", lw=1)
    ax_t.plot(t, T, color="r", lw=1)
    ax_k.plot(t, k, color="b", lw=1)
    (line1,) = ax_cnitrate.plot(t, c_true, color="b", lw=1)
    (line2,) = ax_cnitrate.plot(t, c_obs, color="r", lw=0, markersize=1, marker="o")
    miniax.plot(t, c_true, color="b", lw=1)
    miniax.plot(t, c_obs, color="r", lw=0, markersize=1, marker="o")
    ax_cnitrate.legend([line1, line2], ['True Value', 'Observation'], frameon=False, loc='lower left')

    fig.align_ylabels_coords([ax_do, ax_k], -0.1, 0.5)
    # fig.savefig('Visulization/一级反应数据生成.pdf')
    fig.show()
